Prompt Generator for Stable Diffusion#

Note

Install ekorpkit package first.

Set logging level to Warning, if you don’t want to see verbose logging.

If you run this notebook in Colab, set Hardware accelerator to GPU.

%%capture

!pip install ekorpkit[art]

exit()

Preparing the environment#

%config InlineBackend.figure_format='retina'
%load_ext autotime

from ekorpkit import eKonf

eKonf.setLogger("WARNING")
eKonf.set_cuda(device=0)
print("version:", eKonf.__version__)

is_colab = eKonf.is_colab()
print("is colab?", is_colab)
if is_colab:
    eKonf.mount_google_drive()
project_dir = eKonf.set_workspace(workspace="/content/drive/MyDrive/workspace/", project="ekorpkit-book")
print("project_dir:", project_dir)
INFO:ekorpkit.base:Setting cuda device to ['A100-SXM4-40GB (id:0)']
INFO:ekorpkit.base:Setting EKORPKIT_WORKSPACE_ROOT to /content/drive/MyDrive/workspace/
INFO:ekorpkit.base:Setting EKORPKIT_PROJECT to ekorpkit-book
INFO:ekorpkit.base:Loaded .env from /workspace/projects/ekorpkit-book/config/.env
version: 0.1.40.post0.dev8
is colab? False
project_dir: /content/drive/MyDrive/workspace/projects/ekorpkit-book
time: 3.27 s (started: 2022-11-10 07:42:57 +00:00)

Load a Generator and Generate Prompts#

To download a certain dataset or model checkpoint, you may need to provide a HuggingFace API token. You can get one from here.


import os

# Set HuggingFace API token
os.environ["HF_USER_ACCESS_TOKEN"] = "YOUR_TOKEN"
from ekorpkit.models.art.prompt import PromptGenerator

pgen = PromptGenerator(verbose=False)
2022-11-10 07:43:16.322453: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
time: 6.94 s (started: 2022-11-10 07:43:15 +00:00)

Loading a model#

Before loading a model, you need to train a model first. To train a model, refer to the Training a Generator section.

# ekorpkit/stable-prompts
pgen.load_model(model_name="ekorpkit/stable-prompts")
time: 4.97 s (started: 2022-11-10 07:43:28 +00:00)

Generating prompts#

You can generate prompts using the generate_prompts function. The function takes the following arguments:

  • prompt: The prompt to be used for generating prompts. If None, the prompt will be generated automatically.

  • num_prompts_to_generate: The number of prompts to be generated.

  • generate_images: Whether to generate images for the prompts.

  • num_samples: The number of images to be generated for each prompt.

  • For other arguments, refer to the following code.

eKonf.print(pgen.generate_config)
{'check_start_token': True,
 'max_prompt_length': 50,
 'min_prompt_length': 30,
 'num_prompts_to_generate': 10,
 'prompt': None,
 'temperature': 1.2,
 'top_k': 70,
 'top_p': 0.9}
time: 1.18 ms (started: 2022-11-10 07:43:34 +00:00)
prompts = pgen.generate_prompts(
    prompt="people looking out a lonely city street",
    num_prompts_to_generate=3,
    generate_images=True,
    num_samples=2,
)
Prompt: people looking out a lonely city street in the snow, concept art portrait of a blonde anime boy with blond hair, high detail, digital art, by greg rutkowski
../../../_images/prompt-generator_12_1.png
Prompt: people looking out a lonely city street, artgerm and greg rutkowski and alphonse mucha, a single glowing white male wolf with blue fur, fur, wolf face, ultra realistic, concept art, intricate details
../../../_images/prompt-generator_12_3.png
Prompt: people looking out a lonely city street, trending on artstation by Artgerm and Greg Rutkowski and Alphonse Mucha, 8k, HD
../../../_images/prompt-generator_12_5.png
time: 1min 6s (started: 2022-11-10 07:49:23 +00:00)

Generating images for prompts#

results = pgen.generate_images(
    prompts=prompts,
    num_samples=3,
)
Prompt: people looking out a lonely city street corner. cityscape. By Greg Rutkowski, Ilya Kuvshinov, WLOP, Stanley Artgerm Lau, Ruan Jia and Fenghua Zhong, trending on ArtStation
../../../_images/prompt-generator_14_1.png
Prompt: people looking out a lonely city street below a lake on a white background, artstation, concept art, matte painting by craig mullins and Anato Finnstark, high detail, volumetric lighting, octane render, 4K resolution
../../../_images/prompt-generator_14_3.png
Prompt: people looking out a lonely city street at sunset, fantasy art by greg rutkowski and craig mullins, trending on artstation
../../../_images/prompt-generator_14_5.png
time: 1min 36s (started: 2022-11-10 07:46:05 +00:00)

Generating images for one prompt#

results = pgen.imagine(
    text_prompts=prompts[0],
    num_samples=6,
    num_inference_steps=150,
    guidance_scale=10,
)
Prompt: people looking out a lonely city street in the snow, concept art portrait of a blonde anime boy with blond hair, high detail, digital art, by greg rutkowski
../../../_images/prompt-generator_16_1.png
time: 1min 30s (started: 2022-11-10 07:51:57 +00:00)

Training a Generator#

Preparing a dataset#

You can use any dataset you want. However, the dataset should be in the format of HuggingFace Datasets.

  1. Using a dataset from HuggingFace Hub

dataset = pgen.load_dataset("Gustavosta/Stable-Diffusion-Prompts")
dataset
WARNING:datasets.builder:Using custom data configuration Gustavosta--Stable-Diffusion-Prompts-d22aeec0ba2a9fdb
WARNING:datasets.builder:Reusing dataset parquet (/content/drive/MyDrive/workspace/projects/ekorpkit-book/disco-imagen/cache/Gustavosta___parquet/Gustavosta--Stable-Diffusion-Prompts-d22aeec0ba2a9fdb/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)
DatasetDict({
    train: Dataset({
        features: ['Prompt'],
        num_rows: 73718
    })
    test: Dataset({
        features: ['Prompt'],
        num_rows: 8192
    })
})
time: 4.64 s (started: 2022-11-10 07:13:51 +00:00)
pgen.trainer_config.num_train_epochs = 2
pgen.train(
    model_name="ekorpkit/stable-prompts",
    dataset=dataset,
    prompt_column="Prompt",
    split="train",
)
PyTorch: setting up devices
loading configuration file https://huggingface.co/distilgpt2/resolve/main/config.json from cache at /content/drive/MyDrive/workspace/projects/ekorpkit-book/disco-imagen/cache/f985248d2791fcff97732e4ee263617adec1edb5429a2b8421734c6d14e39bee.422318838d1ec4e061efb4ea29671cb2a044e244dc69229682bebd7cacc81631
Model config GPT2Config {
  "_name_or_path": "distilgpt2",
  "_num_labels": 1,
  "activation_function": "gelu_new",
  "architectures": [
    "GPT2LMHeadModel"
  ],
  "attn_pdrop": 0.1,
  "bos_token_id": 50256,
  "embd_pdrop": 0.1,
  "eos_token_id": 50256,
  "id2label": {
    "0": "LABEL_0"
  },
  "initializer_range": 0.02,
  "label2id": {
    "LABEL_0": 0
  },
  "layer_norm_epsilon": 1e-05,
  "model_type": "gpt2",
  "n_ctx": 1024,
  "n_embd": 768,
  "n_head": 12,
  "n_inner": null,
  "n_layer": 6,
  "n_positions": 1024,
  "reorder_and_upcast_attn": false,
  "resid_pdrop": 0.1,
  "scale_attn_by_inverse_layer_idx": false,
  "scale_attn_weights": true,
  "summary_activation": null,
  "summary_first_dropout": 0.1,
  "summary_proj_to_labels": true,
  "summary_type": "cls_index",
  "summary_use_proj": true,
  "task_specific_params": {
    "text-generation": {
      "do_sample": true,
      "max_length": 50
    }
  },
  "transformers_version": "4.21.2",
  "use_cache": true,
  "vocab_size": 50257
}

loading weights file https://huggingface.co/distilgpt2/resolve/main/pytorch_model.bin from cache at /content/drive/MyDrive/workspace/projects/ekorpkit-book/disco-imagen/cache/43a212e83e76bcb07f45be584cf100676bdbbbe9c13f9e5c1c050049143a832f.a83d881ec4d624fd4b5826dd026e315246c48c67504ff91c0500570e291a54ba
All model checkpoint weights were used when initializing GPT2LMHeadModel.

All the weights of GPT2LMHeadModel were initialized from the model checkpoint at distilgpt2.
If your task is similar to the task the model of the checkpoint was trained on, you can already use GPT2LMHeadModel for predictions without further training.
***** Running training *****
  Num examples = 18847
  Num Epochs = 2
  Instantaneous batch size per device = 1
  Total train batch size (w. parallel, distributed & accumulation) = 8
  Gradient Accumulation steps = 1
  Total optimization steps = 4712
Automatic Weights & Biases logging enabled, to disable set os.environ["WANDB_DISABLED"] = "true"
wandb version 0.13.5 is available! To upgrade, please run: $ pip install wandb --upgrade
Tracking run with wandb version 0.13.3
Run data is saved locally in /workspace/projects/ekorpkit-book/ekorpkit-book/docs/lectures/aiart/wandb/run-20221110_071621-38wtdfq3
[4712/4712 09:59, Epoch 2/2]
Step Training Loss
200 3.800000
400 2.698200
600 2.500200
800 2.353500
1000 2.253200
1200 2.165000
1400 2.111500
1600 2.049800
1800 2.008200
2000 1.974800
2200 1.924800
2400 1.912100
2600 1.869900
2800 1.851300
3000 1.832200
3200 1.807700
3400 1.796100
3600 1.775500
3800 1.767700
4000 1.761300
4200 1.749900
4400 1.735000
4600 1.750500

Training completed. Do not forget to share your model on huggingface.co/models =)
Waiting for W&B process to finish... (success).

Run history:


train/epoch▁▁▂▂▂▃▃▃▃▄▄▄▅▅▅▆▆▆▇▇▇███
train/global_step▁▁▂▂▂▃▃▃▃▄▄▄▅▅▅▆▆▆▇▇▇███
train/learning_rate██▇▇▇▆▆▆▅▅▅▄▄▄▄▃▃▃▂▂▂▁▁
train/loss█▄▄▃▃▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁
train/total_flos
train/train_loss
train/train_runtime
train/train_samples_per_second
train/train_steps_per_second

Run summary:


train/epoch2.0
train/global_step4712
train/learning_rate0.0
train/loss1.7505
train/total_flos9849319716814848.0
train/train_loss2.05561
train/train_runtime606.0063
train/train_samples_per_second62.201
train/train_steps_per_second7.775

Synced ekorpkit/stable-prompts: https://wandb.ai/entelecheia/ekorpkit-book-prompt-generator/runs/38wtdfq3
Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)
Find logs at: ./wandb/run-20221110_071621-38wtdfq3/logs
Configuration saved in /content/drive/MyDrive/workspace/projects/ekorpkit-book/disco-imagen/models/ekorpkit/stable-prompts/config.json
Model weights saved in /content/drive/MyDrive/workspace/projects/ekorpkit-book/disco-imagen/models/ekorpkit/stable-prompts/pytorch_model.bin
time: 11min 13s (started: 2022-11-10 07:15:21 +00:00)
  1. Using a dataset from a text file

prompt_uri = "https://raw.githubusercontent.com/entelecheia/ekorpkit-book/main/assets/data/prompt_parrot.txt"
dataset = pgen.load_dataset("text", data_files=prompt_uri)
dataset
WARNING:datasets.builder:Using custom data configuration default-681f997af4470be8
WARNING:datasets.builder:Reusing dataset text (/content/drive/MyDrive/workspace/projects/ekorpkit-book/disco-imagen/cache/text/default-681f997af4470be8/0.0.0/21a506d1b2b34316b1e82d0bd79066905d846e5d7e619823c0dd338d6f1fa6ad)
DatasetDict({
    train: Dataset({
        features: ['text'],
        num_rows: 185
    })
})
time: 1.29 s (started: 2022-11-10 07:11:28 +00:00)
eKonf.print(pgen.train_config)
eKonf.print(pgen.trainer_config)
{'pretrained_model_name_or_path': 'distilgpt2',
 'prompt_column': 'Prompt',
 'split': 'train'}
{'logging_steps': 200,
 'num_train_epochs': 1,
 'overwrite_output_dir': True,
 'per_device_train_batch_size': 1,
 'prediction_loss_only': True,
 'report_to': 'wandb',
 'run_name': 'ekorpkit/stable-prompts',
 'save_steps': 0,
 'seed': None}
time: 3.54 ms (started: 2022-11-10 07:11:32 +00:00)
pgen.trainer_config.num_train_epochs = 2
pgen.train(model_name="ekorpkit/prompt_parrot", dataset=dataset, prompt_column="text")
Token indices sequence length is longer than the specified maximum sequence length for this model (22608 > 1024). Running this sequence through the model will result in indexing errors
***** Running training *****
  Num examples = 22
  Num Epochs = 2
  Instantaneous batch size per device = 1
  Total train batch size (w. parallel, distributed & accumulation) = 8
  Gradient Accumulation steps = 1
  Total optimization steps = 6
Automatic Weights & Biases logging enabled, to disable set os.environ["WANDB_DISABLED"] = "true"
ERROR:wandb.jupyter:Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.
wandb: Currently logged in as: entelecheia. Use `wandb login --relogin` to force relogin
wandb version 0.13.5 is available! To upgrade, please run: $ pip install wandb --upgrade
Tracking run with wandb version 0.13.3
Run data is saved locally in /workspace/projects/ekorpkit-book/ekorpkit-book/docs/lectures/aiart/wandb/run-20221110_071154-29lg6d7y
[6/6 00:12, Epoch 2/2]
Step Training Loss

Training completed. Do not forget to share your model on huggingface.co/models =)
Waiting for W&B process to finish... (success).

Run history:


train/epoch
train/global_step
train/total_flos
train/train_loss
train/train_runtime
train/train_samples_per_second
train/train_steps_per_second

Run summary:


train/epoch2.0
train/global_step6
train/total_flos11497057026048.0
train/train_loss17.0794
train/train_runtime56.6702
train/train_samples_per_second0.776
train/train_steps_per_second0.106

Synced ekorpkit/prompt_parrot: https://wandb.ai/entelecheia/ekorpkit-book-prompt-generator/runs/29lg6d7y
Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)
Find logs at: ./wandb/run-20221110_071154-29lg6d7y/logs
Configuration saved in /content/drive/MyDrive/workspace/projects/ekorpkit-book/disco-imagen/models/ekorpkit/prompt_parrot/config.json
Model weights saved in /content/drive/MyDrive/workspace/projects/ekorpkit-book/disco-imagen/models/ekorpkit/prompt_parrot/pytorch_model.bin
time: 1min 18s (started: 2022-11-10 07:11:37 +00:00)

References#